package etcd_sdr

import (
	"context"
	"etcd-sdr/pkg/maputils"
	"fmt"
	jsoniter "github.com/json-iterator/go"
	clientv3 "go.etcd.io/etcd/client/v3"
	"go.uber.org/zap"
	"google.golang.org/grpc/resolver"
	"time"
)

type etcdResolver struct {
	namespace  *string
	ctx        context.Context
	cancelFunc context.CancelFunc

	client        *clientv3.Client
	target        resolver.Target
	cc            resolver.ClientConn
	opts          resolver.BuildOptions
	store         *maputils.Map[string, *ServerInfo] // map[serverInfoKey,info]
	ignores       *maputils.Map[string, struct{}]    // map[serverInfoKey,struct{}]
	refreshTicker *time.Ticker
	log           *zap.Logger
}

func (r *etcdResolver) ResolveNow(rn resolver.ResolveNowOptions) {
	if err := r.updateState(); err != nil {
		r.log.Error("[-] etcd resolver update state err")
		return
	}
}

func (r *etcdResolver) Close() {
	r.store = nil
	r.cancelFunc()
}

// todo 排除本机.....
func (r *etcdResolver) updateState() error {
	addresses := make([]resolver.Address, 0)

	// append host address in addresses
	r.store.Range(func(key string, value *ServerInfo) bool {
		if _, ok := r.ignores.Load(key); ok {
			// ignore server
			return true
		}
		addresses = append(addresses, resolver.Address{Addr: value.Target.URL.Host})
		return true
	})

	// update state
	if err := r.cc.UpdateState(resolver.State{Addresses: addresses}); err != nil {
		return err
	}
	return nil
}

func (r *etcdResolver) watch() {

	watchChan := r.client.Watch(r.ctx, getNamingPrefixKey(r.namespace, r.target), clientv3.WithPrefix())

	for {
		select {
		// watch chan
		case resp := <-watchChan:
			if resp.Err() != nil {
				r.log.Error(fmt.Sprintf("[-] etcd resolver watch err: %s", resp.Err()))
				return
			}

			for i := range resp.Events {
				switch resp.Events[i].Type {
				case clientv3.EventTypePut:
					info, err := UnMarshalServerInfo(resp.Events[i].Kv.Value)
					if err != nil {
						return
					}
					r.store.Store(info.GetServerInfoKey(), info)

				case clientv3.EventTypeDelete:
					info, err := UnMarshalServerInfo(resp.Events[i].Kv.Value)
					if err != nil {
						return
					}
					r.store.Delete(info.GetServerInfoKey())

				default:
					r.log.Error("[-] etcd resolver watch type err")
					return
				}
			}

			r.ResolveNow(resolver.ResolveNowOptions{})

		// mandatory full update guarantee
		case <-r.refreshTicker.C:
			r.ResolveNow(resolver.ResolveNowOptions{})

		case <-r.ctx.Done():
			return
		}
	}
}

func newEtcdResolver(
	namespace *string,
	ctx context.Context,
	client *clientv3.Client,
	refreshTicker *time.Ticker,
	target resolver.Target,
	cc resolver.ClientConn,
	opts resolver.BuildOptions,
	log *zap.Logger,
	ignores *maputils.Map[string, struct{}],
) (*etcdResolver, error) {
	_ctx, cancelFunc := context.WithCancel(ctx)

	r := &etcdResolver{
		namespace:     namespace,
		ctx:           _ctx,
		cancelFunc:    cancelFunc,
		client:        client,
		target:        target,
		cc:            cc,
		opts:          opts,
		refreshTicker: refreshTicker,
		log:           log,
		ignores:       ignores,
	}

	r.store = maputils.New[string, *ServerInfo]()
	getResp, err := r.client.Get(r.ctx, getNamingPrefixKey(r.namespace, target), clientv3.WithPrefix())
	if err != nil {
		return nil, err
	}

	for i := range getResp.Kvs {
		var info *ServerInfo
		_ = jsoniter.Unmarshal(getResp.Kvs[i].Value, &info)
		//info, err = UnMarshalServerInfo(getResp.Kvs[i].Value)
		//
		//if err != nil {
		//	return nil, err
		//}
		r.store.Store(info.GetServerInfoKey(), info)
	}

	if err = r.updateState(); err != nil {
		return nil, err
	}
	return r, nil
}
